//
// Created by whitby on 8/4/22.
//

#include <thread>
#include <cmath>

#include "fusion_functions.h"

using namespace std;

void FusionFunctions::initialize(int _width, int _height, float _fx, float _fy, float _cx, float _cy, float _fuse_far,
                                 float _fuse_near) {
    image_width = _width;
    image_height = _height;
    sp_width = image_width / SP_SIZE;
    sp_height = image_height / SP_SIZE;
    fx = _fx;
    fy = _fy;
    cx = _cx;
    cy = _cy;

    fuse_far = _fuse_far;
    fuse_near = _fuse_near;

    superpixel_seeds.resize(sp_width * sp_height);
    superpixel_index.resize(image_width * image_height);
    space_map.resize(image_width * image_height * 3);
    norm_map.resize(image_width * image_height * 3);
}

void FusionFunctions::fuse_initialize_map(int reference_frame_index, cv::Mat &input_image, cv::Mat &input_depth,
                                          Eigen::Matrix4f &pose, vector<shared_ptr<SurfelElement>> &local_surfels,
                                          vector<shared_ptr<SurfelElement>> &new_surfels) {
    image = input_image;
    depth = input_depth;
    local_surfels_ptr = &local_surfels;
    new_surfels_ptr = &new_surfels;

    // generate super-pixels
    generate_super_pixels();

    cout << "generate " << superpixel_seeds.size()  << " super pixels" << endl;

    // fuse
    Eigen::Matrix4f inv_pose = pose.inverse();
    vector<thread> thread_pool;
    for (int i = 0; i < THREAD_NUM; ++i) {
        thread this_thread(&FusionFunctions::fuse_surfels_kernel, this, i, THREAD_NUM,
                           reference_frame_index, pose, inv_pose);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto &thread: thread_pool) {
        if (thread.joinable())
            thread.join();
    }

    // initialize
    initialize_surfels(reference_frame_index, pose);
}

/* generate super pixels */
void FusionFunctions::generate_super_pixels() {
    memset(superpixel_seeds.data(), 0, superpixel_seeds.size() * sizeof(Superpixel_seed));
    std::fill(superpixel_index.begin(), superpixel_index.end(), 0);
    std::fill(norm_map.begin(), norm_map.end(), 0);

    initialize_seeds();

    for (int it_i = 0; it_i < ITERATION_NUM; ++it_i) {
        update_pixels();
        update_seeds();
    }

    calculate_norms();
}

void FusionFunctions::initialize_seeds() {
    vector<thread> thread_pool;
    for (int i = 0; i < THREAD_NUM; i++) {
        thread this_thread(&FusionFunctions::initialize_seeds_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto &thread: thread_pool) {
        if (thread.joinable())
            thread.join();
    }
}

void FusionFunctions::initialize_seeds_kernel(int thread_i, int thread_num) {
    int step = static_cast<int>(superpixel_seeds.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(superpixel_seeds.size());

    for (int seed_i = begin_index; seed_i < end_index; seed_i++) {
        int sp_x = seed_i % sp_width;
        int sp_y = seed_i / sp_width;
        int image_x = sp_x * SP_SIZE + SP_SIZE / 2;
        int image_y = sp_y * SP_SIZE + SP_SIZE / 2;
        image_x = image_x < (image_width - 1) ? image_x : (image_width - 1);
        image_y = image_y < (image_height - 1) ? image_y : (image_height - 1);

        Superpixel_seed this_sp{};
        this_sp.x = static_cast<float>(image_x);
        this_sp.y = static_cast<float>(image_y);
        this_sp.mean_intensity = image.at<uchar>(image_y, image_x);
        this_sp.fused = false;
        this_sp.stable = false;
        this_sp.mean_depth = depth.at<float>(image_y, image_x);

        if (this_sp.mean_depth < 0.01) {  // depth is invalid
            int check_x_begin = sp_x * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
            int check_y_begin = sp_y * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
            int check_x_end = check_x_begin + SP_SIZE * 2;
            int check_y_end = check_y_begin + SP_SIZE * 2;
            check_x_begin = check_x_begin > 0 ? check_x_begin : 0;
            check_y_begin = check_y_begin > 0 ? check_y_begin : 0;
            check_x_end = check_x_end < image_width - 1 ? check_x_end : image_width - 1;
            check_y_end = check_y_end < image_height - 1 ? check_y_end : image_height - 1;
            bool find_depth = false;
            for (int check_j = check_y_begin; check_j < check_y_end; check_j++) {
                for (int check_i = check_x_begin; check_i < check_x_end; check_i++) {
                    float this_depth = depth.at<float>(check_j, check_i);
                    if (this_depth > 0.01) {
                        this_sp.mean_depth = this_depth;
                        find_depth = true;
                        break;
                    }
                }
                if (find_depth)
                    break;
            }
        }
        superpixel_seeds[seed_i] = this_sp;
    }
}

void FusionFunctions::update_pixels() {
    vector<thread> thread_pool;
    for (int i = 0; i < THREAD_NUM; i++) {
        thread this_thread(&FusionFunctions::update_pixels_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto &thread: thread_pool) {
        if (thread.joinable())
            thread.join();
    }
}

void FusionFunctions::update_pixels_kernel(int thread_i, int thread_num) {
    int step_row = image_height / thread_num;
    int start_row = thread_i * step_row;
    int end_row = start_row + step_row;
    if (thread_i == thread_num - 1)
        end_row = image_height;

    for (int row_i = start_row; row_i < end_row; ++row_i) {
        for (int col_i = 0; col_i < image_width; ++col_i) {
            if (superpixel_seeds[superpixel_index[row_i * image_width + col_i]].stable)
                continue;
            float my_intensity = image.at<uchar>(row_i, col_i);
            float my_inv_depth = 0.0;
            if (depth.at<float>(row_i, col_i) > 0.01)
                my_inv_depth = 1.0f / depth.at<float>(row_i, col_i);
            int base_sp_x = col_i / SP_SIZE;
            int base_sp_y = row_i / SP_SIZE;
            float min_dist_depth = 1e6, min_dist_no_depth = 1e6;
            int min_sp_index_depth = -1, min_sp_index_no_depth = -1;
            bool all_has_depth = true;
            for (int check_i = -1; check_i <= 1; ++check_i)
                for (int check_j = -1; check_j <= 1; ++check_j) {
                    int check_sp_x = base_sp_x + check_i;
                    int check_sp_y = base_sp_y + check_j;
                    int dist_sp_x = abs(check_sp_x * SP_SIZE + SP_SIZE / 2 - col_i);
                    int dist_sp_y = abs(check_sp_y * SP_SIZE + SP_SIZE / 2 - row_i);
                    if (dist_sp_x < SP_SIZE && dist_sp_y < SP_SIZE &&
                        check_sp_x >= 0 && check_sp_x < sp_width &&
                        check_sp_y >= 0 && check_sp_y < sp_height) {
                        float dist_depth, dist_no_depth;
                        all_has_depth &= calculate_cost(dist_no_depth, dist_depth,
                                                        my_intensity, my_inv_depth,
                                                        col_i, row_i, check_sp_x, check_sp_y);
                        if (dist_depth < min_dist_depth) {
                            min_dist_depth = dist_depth;
                            min_sp_index_depth = check_sp_y * sp_width + check_sp_x;
                        }
                        if (dist_no_depth < min_dist_no_depth) {
                            min_dist_no_depth = dist_no_depth;
                            min_sp_index_no_depth = check_sp_y * sp_width + check_sp_x;
                        }
                    }
                }
            if (all_has_depth) {
                superpixel_index[row_i * image_width + col_i] = min_sp_index_depth;
                superpixel_seeds[min_sp_index_depth].stable = false;
            } else {
                superpixel_index[row_i * image_width + col_i] = min_sp_index_no_depth;
                superpixel_seeds[min_sp_index_no_depth].stable = false;
            }
        }
    }
}

bool FusionFunctions::calculate_cost(float &no_depth_cost, float &depth_cost, const float &pixel_intensity,
                                     const float &pixel_inverse_depth, const int &x, const int &y, const int &sp_x,
                                     const int &sp_y) {
    int sp_index = sp_y * sp_width + sp_x;
    no_depth_cost = 0;
    float dist = (superpixel_seeds[sp_index].x - static_cast<float>(x)) *
                 (superpixel_seeds[sp_index].x - static_cast<float>(x)) +
                 (superpixel_seeds[sp_index].y - static_cast<float>(y)) *
                 (superpixel_seeds[sp_index].y - static_cast<float>(y));
    auto norm_dist = static_cast<float>(SP_SIZE) / 2.f;
    no_depth_cost += dist / (norm_dist * norm_dist);
    float intensity_diff = (superpixel_seeds[sp_index].mean_intensity - pixel_intensity);
    no_depth_cost += intensity_diff * intensity_diff / 100.0f;
    depth_cost = no_depth_cost;
    if (superpixel_seeds[sp_index].mean_depth > 0 && pixel_inverse_depth > 0) {
        float inverse_depth_diff = 1.0f / superpixel_seeds[sp_index].mean_depth - pixel_inverse_depth;
        depth_cost += inverse_depth_diff * inverse_depth_diff * 400.0f;
        return true;
    }
    return false;
}

void FusionFunctions::update_seeds() {
    vector<thread> thread_pool;
    for (int i = 0; i < THREAD_NUM; ++i) {
        thread this_thread(&FusionFunctions::update_seeds_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto &thread: thread_pool) {
        if (thread.joinable())
            thread.join();
    }
}

void FusionFunctions::update_seeds_kernel(int thread_i, int thread_num) {
    int step = static_cast<int>(superpixel_seeds.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(superpixel_seeds.size());

    for (int seed_i = begin_index; seed_i < end_index; ++seed_i) {
        if (superpixel_seeds[seed_i].stable)
            continue;
        int sp_x = seed_i % sp_width;
        int sp_y = seed_i / sp_width;
        int check_x_begin = sp_x * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
        int check_y_begin = sp_y * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
        int check_x_end = check_x_begin + SP_SIZE * 2;
        int check_y_end = check_y_begin + SP_SIZE * 2;
        check_x_begin = check_x_begin > 0 ? check_x_begin : 0;
        check_y_begin = check_y_begin > 0 ? check_y_begin : 0;
        check_x_end = check_x_end < image_width - 1 ? check_x_end : image_width - 1;
        check_y_end = check_y_end < image_height - 1 ? check_y_end : image_height - 1;

        int sum_x = 0;
        int sum_y = 0;
        int sum_intensity = 0.0;
        int sum_intensity_num = 0.0;
        float sum_depth = 0.0;
        float sum_depth_num = 0.0;
        vector<float> depth_vector;
        for (int check_j = check_y_begin; check_j < check_y_end; check_j++) {
            for (int check_i = check_x_begin; check_i < check_x_end; check_i++) {
                int pixel_index = check_j * image_width + check_i;
                if (superpixel_index[pixel_index] == seed_i) {
                    sum_x += check_i;
                    sum_y += check_j;
                    sum_intensity_num += 1.0;
                    sum_intensity += image.at<uchar>(check_j, check_i);
                    float check_depth = depth.at<float>(check_j, check_i);
                    if (check_depth > 0.01) {
                        depth_vector.push_back(check_depth);
                        sum_depth += check_depth;
                        sum_depth_num += 1.0;
                    }
                }
            }
        }

        if (sum_intensity_num == 0)
            return;

        float pre_intensity = superpixel_seeds[seed_i].mean_intensity;
        float pre_x = superpixel_seeds[seed_i].x;
        float pre_y = superpixel_seeds[seed_i].y;

        superpixel_seeds[seed_i].mean_intensity = static_cast<float>(sum_intensity) / static_cast<float>(sum_intensity_num);
        superpixel_seeds[seed_i].x = static_cast<float>(sum_x) / static_cast<float>(sum_intensity_num);
        superpixel_seeds[seed_i].y = static_cast<float>(sum_y) / static_cast<float>(sum_intensity_num);

        float update_diff = fabs(pre_intensity - superpixel_seeds[seed_i].mean_intensity)
                            + fabs(pre_x - superpixel_seeds[seed_i].x) + fabs(pre_y - superpixel_seeds[seed_i].y);
        if (update_diff < 0.2)
            superpixel_seeds[seed_i].stable = true;

        if (sum_depth_num > 0) {
            float mean_depth = sum_depth / sum_depth_num;
            float sum_a, sum_b;
            for (int newton_i = 0; newton_i < 5; ++newton_i) {
                sum_a = sum_b = 0;
                for (auto pixel_depth : depth_vector) {
                    float residual = mean_depth - pixel_depth;
                    if (fabs(residual) < HUBER_RANGE) {
                        sum_a += 2 * residual;
                        sum_b += 2;
                    } else {
                        sum_a += residual > 0 ? HUBER_RANGE : -1 * HUBER_RANGE;
                    }
                }
                float delta_depth = -sum_a / (sum_b + 10.0f);
                mean_depth = mean_depth + delta_depth;
                if (delta_depth < 0.01 && delta_depth > -0.01)
                    break;
            }
            superpixel_seeds[seed_i].mean_depth = mean_depth;
        } else {
            superpixel_seeds[seed_i].mean_depth = 0.0;
        }
    }
}

void FusionFunctions::calculate_norms() {
    vector<thread> thread_pool;
    for (int i = 0; i < THREAD_NUM; i++) {
        thread this_thread(&FusionFunctions::calculate_spaces_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto & thread : thread_pool) {
        if (thread.joinable())
            thread.join();
    }
    thread_pool.clear();

    for (int i = 0; i < THREAD_NUM; i++) {
        thread this_thread(&FusionFunctions::calculate_pixels_norms_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto & thread : thread_pool) {
        if (thread.joinable())
            thread.join();
    }
    thread_pool.clear();

    for (int i = 0; i < THREAD_NUM; i++) {
        thread this_thread(&FusionFunctions::calculate_sp_depth_norms_kernel, this, i, THREAD_NUM);
        thread_pool.push_back(std::move(this_thread));
    }
    for (auto & thread : thread_pool) {
        if (thread.joinable())
            thread.join();
    }
}

void FusionFunctions::calculate_spaces_kernel(int thread_i, int thread_num) {
    int step_row = image_height / thread_num;
    int start_row = step_row * thread_i;
    int end_row = start_row + step_row;
    if (thread_i == thread_num - 1)
        end_row = image_height;

    for (int row_i = start_row; row_i < end_row; row_i++)
        for (int col_i = 0; col_i < image_width; col_i++) {
            int my_index = row_i * image_width + col_i;
            float my_depth = depth.at<float>(row_i, col_i);
            float x, y, z;
            back_project(static_cast<float>(col_i), static_cast<float>(row_i), my_depth, x, y, z);
            space_map[my_index * 3] = x;
            space_map[my_index * 3 + 1] = y;
            space_map[my_index * 3 + 2] = z;
        }
}

void FusionFunctions::back_project(const float &u, const float &v, const float &pixel_depth, float &x, float &y, float &z) const {
    x = (u - cx) / fx * pixel_depth;
    y = (v - cy) / fy * pixel_depth;
    z = pixel_depth;
}

void FusionFunctions::project(float &x, float &y, float &z, float &u, float &v) const {
    u = x * fx / z + cx;
    v = y * fy / z + cy;
}

void FusionFunctions::calculate_pixels_norms_kernel(int thread_i, int thread_num) {
    int step_row = image_height / thread_num;
    int start_row = step_row * thread_i;
    start_row = start_row > 1 ? start_row : 1;
    int end_row = start_row + step_row;
    if (thread_i == thread_num - 1)
        end_row = image_height - 1;

    for (int row_i = start_row; row_i < end_row; row_i++)
        for (int col_i = 1; col_i < image_width - 1; col_i++) {
            int my_index = row_i * image_width + col_i;
            float my_x, my_y, my_z;
            my_x = space_map[my_index * 3];
            my_y = space_map[my_index * 3 + 1];
            my_z = space_map[my_index * 3 + 2];
            float right_x, right_y, right_z;
            right_x = space_map[my_index * 3 + 3];
            right_y = space_map[my_index * 3 + 4];
            right_z = space_map[my_index * 3 + 5];
            float down_x, down_y, down_z;
            down_x = space_map[my_index * 3 + image_width * 3];
            down_y = space_map[my_index * 3 + image_width * 3 + 1];
            down_z = space_map[my_index * 3 + image_width * 3 + 2];
            if (my_z < 0.1 || right_z < 0.1 || down_z < 0.1)
                continue;
            right_x = right_x - my_x;
            right_y = right_y - my_y;
            right_z = right_z - my_z;
            down_x = down_x - my_x;
            down_y = down_y - my_y;
            down_z = down_z - my_z;
            float norm_x, norm_y, norm_z, norm_length;
            norm_x = right_y * down_z - right_z * down_y;
            norm_y = right_z * down_x - right_x * down_z;
            norm_z = right_x * down_y - right_y * down_x;
            norm_length = std::sqrt(norm_x * norm_x + norm_y * norm_y + norm_z * norm_z);
            norm_x /= norm_length;
            norm_y /= norm_length;
            norm_z /= norm_length;
            float view_angle = (norm_x * my_x + norm_y * my_y + norm_z * my_z)
                               / std::sqrt(my_x * my_x + my_y * my_y + my_z * my_z);
            if (view_angle > -MAX_ANGLE_COS && view_angle < MAX_ANGLE_COS)
                continue;
            norm_map[my_index * 3] = norm_x;
            norm_map[my_index * 3 + 1] = norm_y;
            norm_map[my_index * 3 + 2] = norm_z;
        }
}

void FusionFunctions::calculate_sp_depth_norms_kernel(int thread_i, int thread_num) {
    int step = static_cast<int>(superpixel_seeds.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(superpixel_seeds.size());

    for (int seed_i = begin_index; seed_i < end_index; ++seed_i) {
        int sp_x = seed_i % sp_width;
        int sp_y = seed_i / sp_width;
        int check_x_begin = sp_x * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
        int check_y_begin = sp_y * SP_SIZE + SP_SIZE / 2 - SP_SIZE;
        int check_x_end = check_x_begin + 2 * SP_SIZE;
        int check_y_end = check_y_begin + 2 * SP_SIZE;

        vector<float> pixel_depth, pixel_norms, pixel_positions, pixel_inlier_positions;
        float valid_depth_num = 0;
        float max_dist = 0;
        for (int check_j = check_y_begin; check_j < check_y_end; ++check_j) {
            for (int check_i = check_x_begin; check_i < check_x_end; ++check_i) {
                int pixel_index = check_j * image_width + check_i;
                if (pixel_index < 0 || pixel_index >= superpixel_index.size())
                    continue;
                if (superpixel_index[pixel_index] == seed_i) {
                    float x_diff = static_cast<float>(check_i) - superpixel_seeds[seed_i].x;
                    float y_diff = static_cast<float>(check_j) - superpixel_seeds[seed_i].y;
                    float dist = x_diff * x_diff + y_diff * y_diff;
                    if (dist > max_dist)
                        max_dist = dist;

                    float my_depth = depth.at<float>(check_j, check_i);
                    if (my_depth > 0.05) {
                        pixel_depth.push_back(my_depth);
                        pixel_norms.push_back(norm_map[pixel_index * 3]);
                        pixel_norms.push_back(norm_map[pixel_index * 3 + 1]);
                        pixel_norms.push_back(norm_map[pixel_index * 3 + 2]);
                        valid_depth_num += 1;
                        pixel_positions.push_back(space_map[pixel_index * 3]);
                        pixel_positions.push_back(space_map[pixel_index * 3 + 1]);
                        pixel_positions.push_back(space_map[pixel_index * 3 + 2]);
                    }
                }
            }
        }

        if (valid_depth_num < 16)
            continue;
        float mean_depth = superpixel_seeds[seed_i].mean_depth;
        float norm_x, norm_y, norm_z, norm_b;
        norm_x = norm_y = norm_z = norm_b = 0.0;
        float inlier_num = 0;
        for (int p_i = 0; p_i < pixel_depth.size(); ++p_i) {
            float residual = mean_depth - pixel_depth[p_i];
            if (residual < HUBER_RANGE && residual > -HUBER_RANGE) {
                norm_x += pixel_norms[p_i * 3];
                norm_y += pixel_norms[p_i * 3 + 1];
                norm_z += pixel_norms[p_i * 3 + 2];
                inlier_num += 1;

                pixel_inlier_positions.push_back(pixel_positions[p_i * 3]);
                pixel_inlier_positions.push_back(pixel_positions[p_i * 3 + 1]);
                pixel_inlier_positions.push_back(pixel_positions[p_i * 3 + 2]);
            }
        }

        if (inlier_num / static_cast<float>(pixel_depth.size()) < 0.8)
            continue;

        float norm_length = sqrt(norm_x * norm_x + norm_y * norm_y + norm_z * norm_z);
        norm_x /= norm_length;
        norm_y /= norm_length;
        norm_z /= norm_length;

        get_huber_norm(norm_x, norm_y, norm_z, norm_b, pixel_inlier_positions);

        float avg_x, avg_y, avg_z;
        back_project(superpixel_seeds[seed_i].x, superpixel_seeds[seed_i].y, mean_depth,
                     avg_x, avg_y, avg_z);

        {
            // make sure (avg_x, avg_y, avg_z) on the surfel
            float k = -1.0f * (avg_x * norm_x + avg_y * norm_y + avg_z * norm_z) - norm_b;
            avg_x += k * norm_x;
            avg_y += k * norm_y;
            avg_z += k * norm_z;
            mean_depth = avg_z;
        }

        float view_cos = -1.0f * (norm_x * avg_x + norm_y * avg_y + norm_z * avg_z) /
                         std::sqrt(avg_x * avg_x + avg_y * avg_y + avg_z * avg_z);

        if (view_cos < 0) {
            view_cos *= -1.0;
            norm_x *= -1.0;
            norm_y *= -1.0;
            norm_z *= -1.0;
        }

        superpixel_seeds[seed_i].norm_x = norm_x;
        superpixel_seeds[seed_i].norm_y = norm_y;
        superpixel_seeds[seed_i].norm_z = norm_z;
        superpixel_seeds[seed_i].pos_x = avg_x;
        superpixel_seeds[seed_i].pos_y = avg_y;
        superpixel_seeds[seed_i].pos_z = avg_z;
        superpixel_seeds[seed_i].mean_depth = mean_depth;
        superpixel_seeds[seed_i].view_cos = view_cos;
        superpixel_seeds[seed_i].size = std::sqrt(max_dist);
    }
}

void FusionFunctions::get_huber_norm(float &nx, float &ny, float &nz, float &nb, std::vector<float> &points) {
    int point_num = static_cast<int>(points.size()) / 3;
    float sum_x, sum_y, sum_z;
    sum_x = sum_y = sum_z = 0.0;
    for (int i = 0; i < point_num; ++i) {
        sum_x += points[i * 3];
        sum_y += points[i * 3 + 1];
        sum_z += points[i * 3 + 2];
    }
    auto point_num_f = static_cast<float>(point_num);
    sum_x /= point_num_f;
    sum_y /= point_num_f;
    sum_z /= point_num_f;
    nb = 0;

    for (int i = 0; i < point_num; ++i) {
        points[i * 3] -= sum_x;
        points[i * 3 + 1] -= sum_y;
        points[i * 3 + 2] -= sum_z;
    }

    for (int gn_i = 0; gn_i < 5; ++gn_i) {
        Eigen::Matrix4f hessian = Eigen::Matrix4f::Zero();
        Eigen::Vector4f jacobian = Eigen::Vector4f::Zero();
        for (int i = 0; i < point_num; i++) {
            float residual = points[i * 3] * nx + points[i * 3 + 1] * ny + points[i * 3 + 2] * nz + nb;
            if (residual < HUBER_RANGE && residual > -1 * HUBER_RANGE) {
                jacobian(0) += 2 * residual * points[i * 3];
                jacobian(1) += 2 * residual * points[i * 3 + 1];
                jacobian(2) += 2 * residual * points[i * 3 + 2];
                jacobian(3) += 2 * residual;
                hessian(0, 0) += 2 * points[i * 3] * points[i * 3];
                hessian(0, 1) += 2 * points[i * 3] * points[i * 3 + 1];
                hessian(0, 2) += 2 * points[i * 3] * points[i * 3 + 2];
                hessian(0, 3) += 2 * points[i * 3];
                hessian(1, 0) += 2 * points[i * 3 + 1] * points[i * 3];
                hessian(1, 1) += 2 * points[i * 3 + 1] * points[i * 3 + 1];
                hessian(1, 2) += 2 * points[i * 3 + 1] * points[i * 3 + 2];
                hessian(1, 3) += 2 * points[i * 3 + 1];
                hessian(2, 0) += 2 * points[i * 3 + 2] * points[i * 3];
                hessian(2, 1) += 2 * points[i * 3 + 2] * points[i * 3 + 1];
                hessian(2, 2) += 2 * points[i * 3 + 2] * points[i * 3 + 2];
                hessian(2, 3) += 2 * points[i * 3 + 2];
                hessian(3, 0) += 2 * points[i * 3];
                hessian(3, 1) += 2 * points[i * 3 + 1];
                hessian(3, 2) += 2 * points[i * 3 + 2];
                hessian(3, 3) += 2;
            } else if (residual >= HUBER_RANGE) {
                jacobian(0) += HUBER_RANGE * points[i * 3];
                jacobian(1) += HUBER_RANGE * points[i * 3 + 1];
                jacobian(2) += HUBER_RANGE * points[i * 3 + 2];
                jacobian(3) += HUBER_RANGE;
            } else if (residual <= -1 * HUBER_RANGE) {
                jacobian(0) += -1 * HUBER_RANGE * points[i * 3];
                jacobian(1) += -1 * HUBER_RANGE * points[i * 3 + 1];
                jacobian(2) += -1 * HUBER_RANGE * points[i * 3 + 2];
                jacobian(3) += -1 * HUBER_RANGE;
            }
        }
        hessian(0, 0) += 5;
        hessian(1, 1) += 5;
        hessian(2, 2) += 5;
        hessian(3, 3) += 5;
        Eigen::Vector4f update_value = hessian.inverse() * jacobian;
        nx -= update_value(0);
        ny -= update_value(1);
        nz -= update_value(2);
        nb -= update_value(3);
    }
    nb = nb - (nx * sum_x + ny * sum_y + nz * sum_z);
    float norm_length = std::sqrt(nx * nx + ny * ny + nz * nz);
    nx /= norm_length;
    ny /= norm_length;
    nz /= norm_length;
    nb /= norm_length;
}

/* fuse surfels */
void FusionFunctions::fuse_surfels_kernel(int thread_i, int thread_num, int reference_frame_index, Eigen::Matrix4f pose,
                                          Eigen::Matrix4f inv_pose) {
    vector<shared_ptr<SurfelElement>> &local_surfels = *local_surfels_ptr;

    int step = static_cast<int>(local_surfels.size()) / thread_num;
    int begin_index = step * thread_i;
    int end_index = begin_index + step;
    if (thread_i == thread_num - 1)
        end_index = static_cast<int>(local_surfels.size());

    for (int i = begin_index; i < end_index; ++i) {
        // remove unstable
        if (reference_frame_index - local_surfels[i]->last_update > 5 && local_surfels[i]->update_times < 5) {
            local_surfels[i]->update_times = 0;
            continue;
        }

        if (local_surfels[i]->update_times == 0)
            continue;

        Eigen::Vector4f surfel_p_w;
        surfel_p_w(0) = local_surfels[i]->px;
        surfel_p_w(1) = local_surfels[i]->py;
        surfel_p_w(2) = local_surfels[i]->pz;
        surfel_p_w(3) = 1.0;
        Eigen::Vector4f surfel_p_c = inv_pose * surfel_p_w;
        if (surfel_p_c(2) < fuse_near || surfel_p_c(2) > fuse_far)
            continue;

        Eigen::Vector3f norm_w;
        norm_w(0) = local_surfels[i]->nx;
        norm_w(1) = local_surfels[i]->ny;
        norm_w(2) = local_surfels[i]->nz;
        Eigen::Vector3f norm_c;
        norm_c = inv_pose.block<3, 3>(0, 0) * norm_w;

        float project_u, project_v;
        project(surfel_p_c(0), surfel_p_c(1), surfel_p_c(2), project_u, project_v);
        int p_u_int = int(round(project_u));
        int p_v_int = int(round(project_v));
        if (p_u_int < 1 || p_u_int > image_width - 2 || p_v_int < 1 || p_v_int > image_height - 2)
            continue;
        if (surfel_p_c(2) < depth.at<float>(p_v_int, p_u_int) - 1.0) {
            local_surfels[i]->update_times = 0;
            continue;
        }

        int sp_index = superpixel_index[p_v_int * image_width + p_u_int];
        if (superpixel_seeds[sp_index].norm_x == 0 && superpixel_seeds[sp_index].norm_y == 0 &&
            superpixel_seeds[sp_index].norm_z == 0)
            continue;
        if (superpixel_seeds[sp_index].view_cos < MAX_ANGLE_COS)
            continue;

        float camera_f = (fabs(fx) + fabs(fy)) / 2.0f;
        float tolerate_diff = surfel_p_c(2) * surfel_p_c(2)
                              / (BASELINE * camera_f) * DISPARITY_ERROR;
        tolerate_diff = tolerate_diff < MIN_TOLERATE_DIFF ? MIN_TOLERATE_DIFF : tolerate_diff;
        if (fabs(surfel_p_c(2) - superpixel_seeds[sp_index].mean_depth) > tolerate_diff)
            continue;

        float norm_diff_cos = norm_c(0) * superpixel_seeds[sp_index].norm_x
                              + norm_c(1) * superpixel_seeds[sp_index].norm_y
                              + norm_c(2) * superpixel_seeds[sp_index].norm_z;
        if (norm_diff_cos < MAX_ANGLE_COS) {
            local_surfels[i]->update_times = 0;
            continue;
        }

        float old_weight = local_surfels[i]->weight;
        float new_weight = get_weight(superpixel_seeds[sp_index].mean_depth);
        float sum_weight = old_weight + new_weight;

        Eigen::Vector4f sp_p_c, sp_p_w;
        sp_p_c(0) = superpixel_seeds[sp_index].pos_x;
        sp_p_c(1) = superpixel_seeds[sp_index].pos_y;
        sp_p_c(2) = superpixel_seeds[sp_index].pos_z;
        sp_p_c(3) = 1.0;
        sp_p_w = pose * sp_p_c;

        float fused_px = (local_surfels[i]->px * old_weight + new_weight * sp_p_w(0)) / sum_weight;
        float fused_py = (local_surfels[i]->py * old_weight + new_weight * sp_p_w(1)) / sum_weight;
        float fused_pz = (local_surfels[i]->pz * old_weight + new_weight * sp_p_w(2)) / sum_weight;
        float fused_nx = norm_c(0) * old_weight + new_weight * superpixel_seeds[sp_index].norm_x;
        float fused_ny = norm_c(1) * old_weight + new_weight * superpixel_seeds[sp_index].norm_y;
        float fused_nz = norm_c(2) * old_weight + new_weight * superpixel_seeds[sp_index].norm_z;
        float new_norm_length = std::sqrt(fused_nx * fused_nx + fused_ny * fused_ny + fused_nz * fused_nz);
        fused_nx /= new_norm_length;
        fused_ny /= new_norm_length;
        fused_nz /= new_norm_length;

        Eigen::Vector3f new_norm_c, new_norm_w;
        new_norm_c(0) = fused_nx;
        new_norm_c(1) = fused_ny;
        new_norm_c(2) = fused_nz;
        new_norm_w = pose.block<3, 3>(0, 0) * new_norm_c;

        local_surfels[i]->px = fused_px;
        local_surfels[i]->py = fused_py;
        local_surfels[i]->pz = fused_pz;
        local_surfels[i]->nx = new_norm_w(0);
        local_surfels[i]->ny = new_norm_w(1);
        local_surfels[i]->nz = new_norm_w(2);
        local_surfels[i]->weight = sum_weight;
        local_surfels[i]->color = superpixel_seeds[sp_index].mean_intensity;

        float new_size = superpixel_seeds[sp_index].size *
                         fabs(superpixel_seeds[sp_index].mean_depth / (camera_f * superpixel_seeds[sp_index].view_cos));
        if (new_size < local_surfels[i]->size)
            local_surfels[i]->size = new_size;
        local_surfels[i]->last_update = reference_frame_index;
        local_surfels[i]->update_times += 1;
        superpixel_seeds[sp_index].fused = true;
    }
}

float FusionFunctions::get_weight(float &_depth) {
    return min(1.0f / _depth / _depth, 1.0f);
}

void FusionFunctions::initialize_surfels(int reference_frame_index, Eigen::Matrix4f pose) {
    vector<shared_ptr<SurfelElement>> &new_surfels = *new_surfels_ptr;
    new_surfels.clear();
    Eigen::Vector4f position_temp_c, position_temp_w;
    Eigen::Vector3f norm_temp_c, norm_temp_w;
    for (auto &super_pixel: superpixel_seeds) {
        if (super_pixel.mean_depth == 0 || super_pixel.fused || super_pixel.view_cos < MAX_ANGLE_COS)
            continue;
        position_temp_c(0) = super_pixel.pos_x;
        position_temp_c(1) = super_pixel.pos_y;
        position_temp_c(2) = super_pixel.pos_z;
        position_temp_c(3) = 1.0;
        norm_temp_c(0) = super_pixel.norm_x;
        norm_temp_c(1) = super_pixel.norm_y;
        norm_temp_c(2) = super_pixel.norm_z;
        if (norm_temp_c(0) == 0 && norm_temp_c(1) == 0 && norm_temp_c(2) == 0)
            continue;
        position_temp_w = pose * position_temp_c;
        norm_temp_w = pose.block<3, 3>(0, 0) * norm_temp_c;

        SurfelElement new_ele{};
        new_ele.px = position_temp_w(0);
        new_ele.py = position_temp_w(1);
        new_ele.pz = position_temp_w(2);
        new_ele.nx = norm_temp_w(0);
        new_ele.ny = norm_temp_w(1);
        new_ele.nz = norm_temp_w(2);
        float camera_f = (fabs(fx) + fabs(fy)) / 2.0f;
        float new_size = super_pixel.size *
                         fabs(super_pixel.mean_depth / (camera_f * super_pixel.view_cos));
        new_ele.size = new_size;
        new_ele.color = super_pixel.mean_intensity;
        new_ele.weight = get_weight(super_pixel.mean_depth);
        new_ele.update_times = 1;
        new_ele.last_update = reference_frame_index;
        new_surfels.push_back(make_shared<SurfelElement>(new_ele));
    }
}